import torch
import torch.nn as nn
import numpy as np

class KNN(nn.Module):
    def __init__(self, k):
        super(KNN, self).__init__()
        self.k = k

    def fit(self, X_train, Y_train):
        self.X_train = X_train
        self.Y_train = Y_train

    def forward(self, X, pb_info_idx, chosen):
        distances = torch.cdist(X, self.X_train)  # 计算查询样本与训练集中每个样本的距离
        _, indices = torch.topk(distances, self.k, largest=False, dim=1)  # 找出距离最近的 k 个样本的索引

        indices = indices[0].tolist()
        temp = []
        for va in indices:
            if pb_info_idx[va] not in chosen:
                temp.append(va)
        indices = torch.tensor(temp).unsqueeze(0).cuda()
        knn_labels = self.Y_train[indices]  # 根据索引获取对应的标签
        predictions, _ = torch.mode(knn_labels, dim=1)  # 对 k 个最近邻样本的标签求众数，即预测标签

        return pb_info_idx[predictions.item()]

